# -*- coding: utf-8 -*-
"""
配置文件 - 定义模型和训练参数
包含模型架构、训练超参数、数据路径等配置
"""

import os
from dataclasses import dataclass
from typing import Optional


@dataclass
class ModelConfig:
    """模型配置类"""
    # 模型架构参数
    vocab_size: int = 50000  # 词汇表大小
    d_model: int = 512  # 模型维度
    n_heads: int = 8  # 注意力头数
    n_layers: int = 6  # Transformer层数
    d_ff: int = 2048  # 前馈网络维度
    max_seq_length: int = 512  # 最大序列长度
    dropout: float = 0.1  # Dropout率
    
    # 位置编码
    use_positional_encoding: bool = True
    
    # 激活函数
    activation: str = "relu"  # relu, gelu, swish
    
    # 层归一化
    layer_norm_eps: float = 1e-6
    
    # 权重初始化
    initializer_range: float = 0.02


@dataclass
class TrainingConfig:
    """训练配置类"""
    # 基础训练参数
    batch_size: int = 32
    learning_rate: float = 5e-4
    num_epochs: int = 10
    warmup_steps: int = 4000
    max_grad_norm: float = 1.0  # 梯度裁剪
    
    # 优化器参数
    optimizer: str = "adamw"  # adam, adamw, sgd
    weight_decay: float = 0.01
    beta1: float = 0.9
    beta2: float = 0.999
    eps: float = 1e-8
    
    # 学习率调度
    scheduler: str = "cosine"  # linear, cosine, constant
    
    # 保存和日志
    save_steps: int = 1000
    logging_steps: int = 100
    eval_steps: int = 500
    
    # 早停
    early_stopping_patience: int = 3
    early_stopping_threshold: float = 0.001
    
    # 混合精度训练
    use_fp16: bool = True
    
    # 梯度累积
    gradient_accumulation_steps: int = 1


@dataclass
class DataConfig:
    """数据配置类"""
    # 数据路径
    train_data_path: str = "data/train.txt"
    val_data_path: str = "data/val.txt"
    test_data_path: str = "data/test.txt"
    
    # 词汇表路径
    vocab_path: str = "data/vocab.txt"
    
    # 数据预处理
    max_length: int = 512
    min_length: int = 10
    
    # 数据加载
    num_workers: int = 4
    pin_memory: bool = True
    
    # 文本处理
    lowercase: bool = True
    remove_punctuation: bool = False
    
    # 特殊标记
    pad_token: str = "<pad>"
    unk_token: str = "<unk>"
    bos_token: str = "<bos>"
    eos_token: str = "<eos>"
    mask_token: str = "<mask>"


@dataclass
class SystemConfig:
    """系统配置类"""
    # 设备配置
    device: str = "cuda"  # cuda, cpu, auto
    num_gpus: int = 1
    
    # 随机种子
    seed: int = 42
    
    # 输出路径
    output_dir: str = "outputs"
    model_save_dir: str = "models"
    log_dir: str = "logs"
    
    # 检查点
    resume_from_checkpoint: Optional[str] = None
    
    # 实验名称
    experiment_name: str = "transformer_lm"
    
    # 是否使用wandb
    use_wandb: bool = False
    wandb_project: str = "transformer-language-model"


class Config:
    """主配置类，整合所有配置"""
    
    def __init__(self):
        self.model = ModelConfig()
        self.training = TrainingConfig()
        self.data = DataConfig()
        self.system = SystemConfig()
        
        # 创建必要的目录
        self._create_directories()
    
    def _create_directories(self):
        """创建必要的目录"""
        dirs_to_create = [
            self.system.output_dir,
            self.system.model_save_dir,
            self.system.log_dir,
            "data"
        ]
        
        for dir_path in dirs_to_create:
            os.makedirs(dir_path, exist_ok=True)
    
    def update_config(self, **kwargs):
        """更新配置参数"""
        for key, value in kwargs.items():
            if hasattr(self.model, key):
                setattr(self.model, key, value)
            elif hasattr(self.training, key):
                setattr(self.training, key, value)
            elif hasattr(self.data, key):
                setattr(self.data, key, value)
            elif hasattr(self.system, key):
                setattr(self.system, key, value)
            else:
                print(f"Warning: Unknown config parameter: {key}")
    
    def save_config(self, path: str):
        """保存配置到文件"""
        import json
        
        config_dict = {
            "model": self.model.__dict__,
            "training": self.training.__dict__,
            "data": self.data.__dict__,
            "system": self.system.__dict__
        }
        
        with open(path, 'w', encoding='utf-8') as f:
            json.dump(config_dict, f, indent=2, ensure_ascii=False)
    
    def load_config(self, path: str):
        """从文件加载配置"""
        import json
        
        with open(path, 'r', encoding='utf-8') as f:
            config_dict = json.load(f)
        
        # 更新配置
        for section, params in config_dict.items():
            if hasattr(self, section):
                for key, value in params.items():
                    setattr(getattr(self, section), key, value)


# 创建全局配置实例
config = Config()


# 预定义的配置模板
class ConfigTemplates:
    """预定义的配置模板"""
    
    @staticmethod
    def small_model():
        """小型模型配置"""
        cfg = Config()
        cfg.model.d_model = 256
        cfg.model.n_heads = 4
        cfg.model.n_layers = 4
        cfg.model.d_ff = 1024
        cfg.training.batch_size = 64
        return cfg
    
    @staticmethod
    def medium_model():
        """中型模型配置"""
        cfg = Config()
        cfg.model.d_model = 512
        cfg.model.n_heads = 8
        cfg.model.n_layers = 6
        cfg.model.d_ff = 2048
        cfg.training.batch_size = 32
        return cfg
    
    @staticmethod
    def large_model():
        """大型模型配置"""
        cfg = Config()
        cfg.model.d_model = 768
        cfg.model.n_heads = 12
        cfg.model.n_layers = 12
        cfg.model.d_ff = 3072
        cfg.training.batch_size = 16
        return cfg


if __name__ == "__main__":
    # 测试配置
    print("模型配置:")
    print(f"  模型维度: {config.model.d_model}")
    print(f"  注意力头数: {config.model.n_heads}")
    print(f"  层数: {config.model.n_layers}")
    
    print("\n训练配置:")
    print(f"  批次大小: {config.training.batch_size}")
    print(f"  学习率: {config.training.learning_rate}")
    print(f"  训练轮数: {config.training.num_epochs}")
    
    # 保存配置示例
    config.save_config("config_example.json")
    print("\n配置已保存到 config_example.json")